"""
@Project ：utils
@File    ：thread_pool.py
@IDE     ：PyCharm
@Author  ：靓仔
@Date    ：2025/2/13 14:24
@Description ：智能线程池管理器模块
"""
import threading
import concurrent.futures
import multiprocessing
import time
import queue
import heapq
from collections import deque, defaultdict
from threading import Lock, Event
from dataclasses import dataclass
from typing import Callable, Any, Optional, Dict, List
import statistics


@dataclass
class TaskInfo:
    """任务信息"""
    priority: int  # 优先级（数字越小优先级越高）
    task_id: str
    text_length: int
    estimated_time: float
    submit_time: float
    func: Callable
    args: tuple
    kwargs: dict


class IntelligentThreadPool:
    """真正智能的线程池管理器"""
    
    def __init__(self, initial_workers=None, max_workers=None, min_workers=2):
        # 动态线程数配置
        cpu_count = multiprocessing.cpu_count()
        self.min_workers = max(min_workers, 2)
        # 使用更智能的默认值：推荐线程数不超过CPU核心数*2，但不限制在8个
        default_workers = min(max(2, cpu_count * 2), 16)
        self.initial_workers = initial_workers or default_workers
        self.max_workers = max_workers or min(cpu_count * 4, 20)  # 更高上限，支持高并发
        
        # 当前活跃线程数
        self.current_workers = self.initial_workers
        
        # 任务队列和状态
        self.priority_queue = []  # 优先级队列 (priority, task_info)
        self.active_tasks = {}  # task_id -> future
        self.completed_tasks = 0
        self.total_tasks = 0
        self.failed_tasks = 0
        
        # 性能监控
        self.task_completion_times = deque(maxlen=50)  # 最近50个任务的完成时间
        self.network_latencies = deque(maxlen=20)  # 网络延迟历史
        self.worker_utilization = defaultdict(float)  # 线程利用率
        
        # 同步控制
        self.lock = Lock()
        self.task_available = Event()
        self.shutdown_flag = False
        
        # 线程池实例
        self.executor = None
        
        # 自适应调整参数
        self.last_adjustment_time = time.time()
        self.adjustment_interval = 5.0  # 5秒调整一次
        self.performance_window = 10.0  # 10秒性能窗口
        
    def calculate_task_priority(self, text_length: int, task_type: str = "normal") -> int:
        """智能计算任务优先级"""
        base_priority = 100
        
        # 短文本优先处理（减少等待时间）
        if text_length < 200:
            priority = base_priority - 30
        elif text_length < 500:
            priority = base_priority - 10
        elif text_length < 1000:
            priority = base_priority
        else:
            priority = base_priority + 20
            
        # 考虑当前队列状态
        with self.lock:
            queue_length = len(self.priority_queue)
            if queue_length > 10:  # 队列拥堵时，优先处理短任务
                if text_length < 300:
                    priority -= 20
                    
        return priority
    
    def estimate_task_time(self, text_length: int) -> float:
        """基于历史数据估算任务完成时间"""
        if not self.task_completion_times:
            # 没有历史数据时的估算（每100字符约1秒）
            return max(1.0, text_length / 100.0)
            
        # 基于历史数据的智能估算
        avg_time = statistics.mean(self.task_completion_times)
        char_per_second = 100  # 基础处理速度
        
        if len(self.task_completion_times) > 5:
            # 计算实际的字符处理速度
            recent_times = list(self.task_completion_times)[-5:]
            avg_recent = statistics.mean(recent_times)
            char_per_second = max(50, 150 - (avg_recent * 10))
            
        estimated = text_length / char_per_second
        
        # 考虑网络延迟
        if self.network_latencies:
            avg_latency = statistics.mean(self.network_latencies)
            estimated += avg_latency
            
        return max(0.5, estimated)
    
    def should_adjust_workers(self) -> bool:
        """判断是否需要调整线程数"""
        current_time = time.time()
        return (current_time - self.last_adjustment_time) > self.adjustment_interval
    
    def calculate_optimal_workers(self) -> int:
        """计算最优线程数"""
        with self.lock:
            active_count = len(self.active_tasks)
            queue_length = len(self.priority_queue)
            
        # 基础计算：根据任务队列长度
        if queue_length == 0:
            optimal = max(self.min_workers, self.current_workers - 1)
        elif queue_length < 3:
            optimal = self.current_workers
        elif queue_length < 8:
            optimal = min(self.max_workers, self.current_workers + 1)
        else:
            optimal = min(self.max_workers, self.current_workers + 2)
            
        # 性能调整：基于完成时间趋势
        if len(self.task_completion_times) > 10:
            recent_times = list(self.task_completion_times)[-5:]
            older_times = list(self.task_completion_times)[-10:-5]
            
            if recent_times and older_times:
                recent_avg = statistics.mean(recent_times)
                older_avg = statistics.mean(older_times)
                
                # 如果最近任务变慢，增加线程
                if recent_avg > older_avg * 1.2:
                    optimal = min(self.max_workers, optimal + 1)
                # 如果最近任务变快，可以减少线程
                elif recent_avg < older_avg * 0.8:
                    optimal = max(self.min_workers, optimal - 1)
                    
        return optimal
    
    def adjust_thread_pool(self):
        """动态调整线程池大小"""
        if not self.should_adjust_workers():
            return
            
        optimal_workers = self.calculate_optimal_workers()
        
        if optimal_workers != self.current_workers:
            print(f"🔧 智能调整线程数: {self.current_workers} -> {optimal_workers}")
            
            # 创建新的线程池
            old_executor = self.executor
            self.current_workers = optimal_workers
            self.executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=self.current_workers,
                thread_name_prefix="Smart-TTS-Worker"
            )
            
            # 优雅关闭旧线程池
            if old_executor:
                threading.Thread(
                    target=lambda: old_executor.shutdown(wait=True),
                    daemon=True
                ).start()
                
            self.last_adjustment_time = time.time()
    
    def start(self):
        """启动智能线程池"""
        if self.executor is None:
            self.executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=self.current_workers,
                thread_name_prefix="Smart-TTS-Worker"
            )
            print(f"🚀 启动智能线程池: {self.current_workers} 个工作线程")
    
    def shutdown(self):
        """关闭线程池"""
        self.shutdown_flag = True
        if self.executor:
            print("🛑 正在关闭智能线程池...")
            self.executor.shutdown(wait=True)
            self.executor = None
    
    def submit_intelligent_task(self, task_func, text_length: int, task_id: str = None, 
                              task_type: str = "normal", *args, **kwargs):
        """提交智能任务（带优先级和性能监控）"""
        if self.executor is None:
            self.start()
            
        # 动态调整线程池
        self.adjust_thread_pool()
        
        task_id = task_id or f"task_{int(time.time() * 1000)}"
        priority = self.calculate_task_priority(text_length, task_type)
        estimated_time = self.estimate_task_time(text_length)
        
        task_info = TaskInfo(
            priority=priority,
            task_id=task_id,
            text_length=text_length,
            estimated_time=estimated_time,
            submit_time=time.time(),
            func=task_func,
            args=args,
            kwargs=kwargs
        )
        
        with self.lock:
            self.total_tasks += 1
            heapq.heappush(self.priority_queue, (priority, task_info))
        
        # 立即处理任务（不等待队列）
        future = self._process_next_task()
        
        return future
    
    def _process_next_task(self):
        """处理下一个最高优先级任务"""
        with self.lock:
            if not self.priority_queue:
                return None
                
            priority, task_info = heapq.heappop(self.priority_queue)
            
        # 提交任务到线程池
        start_time = time.time()
        future = self.executor.submit(task_info.func, *task_info.args, **task_info.kwargs)
        
        # 添加性能监控回调
        def task_completion_callback(fut):
            end_time = time.time()
            completion_time = end_time - start_time
            
            with self.lock:
                self.completed_tasks += 1
                self.task_completion_times.append(completion_time)
                
                # 移除活跃任务
                if task_info.task_id in self.active_tasks:
                    del self.active_tasks[task_info.task_id]
                    
                # 记录网络延迟（估算）
                if completion_time > 1.0:  # 超过1秒的任务考虑网络因素
                    estimated_network_time = completion_time - task_info.estimated_time
                    if estimated_network_time > 0:
                        self.network_latencies.append(estimated_network_time)
                        
                try:
                    result = fut.result()
                    print(f"✅ 智能任务完成: {task_info.task_id} "
                          f"({task_info.text_length}字符, {completion_time:.2f}秒)")
                except Exception as e:
                    self.failed_tasks += 1
                    print(f"❌ 任务失败: {task_info.task_id} - {str(e)}")
        
        future.add_done_callback(task_completion_callback)
        
        # 记录活跃任务
        with self.lock:
            self.active_tasks[task_info.task_id] = future
            
        return future
    
    def submit_task(self, task_func, *args, **kwargs):
        """兼容旧接口的任务提交"""
        # 尝试从参数中提取文本长度
        text_length = 500  # 默认值
        
        if args and len(args) > 0:
            # 假设第一个参数是任务数据，尝试提取文本长度
            task_data = args[0]
            if isinstance(task_data, (tuple, list)) and len(task_data) > 0:
                text = task_data[0]
                if isinstance(text, str):
                    text_length = len(text)
                    
        return self.submit_intelligent_task(task_func, text_length, *args, **kwargs)
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计信息"""
        with self.lock:
            stats = {
                'current_workers': self.current_workers,
                'active_tasks': len(self.active_tasks),
                'queued_tasks': len(self.priority_queue),
                'completed_tasks': self.completed_tasks,
                'failed_tasks': self.failed_tasks,
                'total_tasks': self.total_tasks,
                'success_rate': ((self.completed_tasks / max(1, self.completed_tasks + self.failed_tasks)) * 100) if (self.completed_tasks + self.failed_tasks) > 0 else 0,
            }
            
            if self.task_completion_times:
                stats.update({
                    'avg_completion_time': statistics.mean(self.task_completion_times),
                    'min_completion_time': min(self.task_completion_times),
                    'max_completion_time': max(self.task_completion_times),
                })
                
            if self.network_latencies:
                stats['avg_network_latency'] = statistics.mean(self.network_latencies)
                
        return stats
    
    def get_progress(self):
        """获取总体进度"""
        with self.lock:
            if self.total_tasks == 0:
                return 0
            return int((self.completed_tasks / self.total_tasks) * 100)
    
    def get_active_tasks(self):
        """获取活跃任务数量"""
        with self.lock:
            return len(self.active_tasks)
    
    def reset_counters(self):
        """重置任务计数器和性能数据"""
        with self.lock:
            self.total_tasks = 0
            self.completed_tasks = 0
            self.failed_tasks = 0
            self.active_tasks.clear()
            self.priority_queue.clear()
            self.task_completion_times.clear()
            self.network_latencies.clear()


# 保持向后兼容
class SmartThreadPool(IntelligentThreadPool):
    """向后兼容的别名"""
    pass
